import numpy as np
import torch

class DynamicWeightAggregator:
    def __init__(self, similarity_threshold=0.8, alpha=0.6, beta=1.5):
        """
        动态权重聚合算法实现 (对应论文算法2)
        
        参数:
            similarity_threshold: 梯度相似性阈值 (θ)
            alpha: 数据质量评估中准确率权重系数
            beta: 相似性放大系数 (γ)
        """
        self.similarity_threshold = similarity_threshold
        self.alpha = alpha
        self.beta = beta
        
    def calculate_weights(self, noisy_grads, global_grad, data_sizes, accuracies):
        """
        计算动态聚合权重
        
        参数:
            noisy_grads: 客户端加噪梯度列表 [grad_client1, grad_client2, ...]
            global_grad: 全局模型梯度
            data_sizes: 客户端数据量列表
            accuracies: 客户端历史准确率列表
            
        返回:
            weights: 归一化的聚合权重
        """
        num_clients = len(noisy_grads)
        weights = np.zeros(num_clients)
        
        # 获取全局梯度向量
        global_grad_vec = self._flatten_grad(global_grad)
        
        for i in range(num_clients):
            # Step 1: 梯度相似性计算 (噪声鲁棒过滤)
            client_grad_vec = self._flatten_grad(noisy_grads[i])
            cos_sim = self._cosine_similarity(client_grad_vec, global_grad_vec)
            
            # 过滤噪声引起的异常梯度
            if cos_sim < self.similarity_threshold:
                weights[i] = 0.0
                continue
                
            # Step 2: 数据质量评估
            # 历史准确率权重 (Sigmoid函数平滑)
            acc_weight = 1 / (1 + np.exp(-accuracies[i]))
            # 数据规模权重
            size_weight = data_sizes[i] / max(data_sizes)
            # 综合质量得分
            quality_score = self.alpha * acc_weight + (1 - self.alpha) * size_weight
            
            # Step 3: 动态权重合成 (指数加权融合)
            weights[i] = np.exp(self.beta * cos_sim) * quality_score
        
        # 归一化权重 (防止除零错误)
        if weights.sum() > 0:
            weights = weights / weights.sum()
        else:
            weights = np.ones(num_clients) / num_clients
            
        return weights
    
    def _flatten_grad(self, grad_dict):
        """将梯度字典展平为向量"""
        return torch.cat([g.flatten() for g in grad_dict.values()])
    
    def _cosine_similarity(self, vec1, vec2):
        """计算余弦相似度"""
        return torch.dot(vec1, vec2) / (torch.norm(vec1) * torch.norm(vec2) + 1e-8)
